import numpy as np
import regression
from matplotlib import pylab as pl

# 定义训练数据
x, y = regression.loadDataSet('data.txt')


# 求取回归方程
b0, b1 = regression.fit(x,y)
print('Line is:y = %2.0fx + %2.0f'%(b0,b1))

# 预测
x_test = np.array([1, 1.5, 2, 3, 4])
y_test = np.zeros((1, len(x_test)))
for i in range(len(x_test)):
    y_test[0][i] = regression.predit(x_test[i], b0, b1)

# 绘制图像
xx = np.linspace(0, 5)
yy = b0*xx + b1
pl.plot(xx, yy, 'k-')
pl.scatter(x, y, cmap=pl.cm.Paired)
pl.scatter(x_test, y_test[0], cmap=pl.cm.Paired)
pl.show()
